/*
 * Copyright 2018 TiKV Project Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package com.webank.tikv;

import com.google.protobuf.ByteString;
import com.webank.tikv.exception.RetryableException;
import com.webank.tikv.kvproto.Kvrpcpb;
import com.webank.tikv.kvproto.Metapb;
import com.webank.tikv.region.RegionStoreClient;
import com.webank.tikv.region.RegionStoreClient.RegionStoreClientBuilder;
import com.webank.tikv.region.TiRegion;
import com.webank.tikv.region.TiStore;
import com.webank.tikv.util.BackOffer;
import com.webank.tikv.util.ConcreteBackOffer;
import org.apache.commons.lang3.tuple.Triple;

import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import static com.webank.tikv.util.BackOffer.REDIS_EXEC_MAX_BACKOFF;

public class RawKVClient {

    private final TiSession tiSession;
    private final RegionStoreClientBuilder clientBuilder;

    public RawKVClient(TiSession session, RegionStoreClientBuilder clientBuilder) {
        Objects.requireNonNull(session, "session is null");
        Objects.requireNonNull(clientBuilder, "clientBuilder is null");
        this.tiSession = session;
        this.clientBuilder = clientBuilder;
    }

    public Kvrpcpb.RedisResponse execRedis(ByteString routeKey, String cmd, ByteString key, ByteString... value) {
        ConcreteBackOffer backOffer = ConcreteBackOffer.newDeadlineBackOff(REDIS_EXEC_MAX_BACKOFF);
        RegionStoreClient client = clientBuilder.build(routeKey, backOffer);
        return client.execRedis(backOffer, cmd, key, value);
    }

    public CompletableFuture<Kvrpcpb.RedisResponse> execRedisAsync(ScheduledExecutorService executor, ByteString routeKey,
                                                                   String cmd, ByteString key, ByteString... value) {
        ConcreteBackOffer backOffer = ConcreteBackOffer.newDeadlineBackOff(REDIS_EXEC_MAX_BACKOFF);
        RegionStoreClient client = clientBuilder.build(routeKey, backOffer);
        return client.execRedisAsync(backOffer, executor, cmd, key, value);
    }

    public CompletableFuture<Kvrpcpb.RedisResponse> directExecRedis(ScheduledExecutorService executor,
                                                                    CompletableFuture<Kvrpcpb.RedisResponse> future,
                                                                    Triple<TiRegion, TiStore, Metapb.Peer> triple,
                                                                    BackOffer backOffer, long sendTime,
                                                                    String cmd, ByteString encodedKey,
                                                                    ByteString... value) {
        TiRegion region = triple.getLeft();
        Metapb.Peer peer = triple.getRight();
        boolean isLeader = region.getLeader().getId() == peer.getId();
        Kvrpcpb.Context context = isLeader ? region.makeLeaderContext() : region.makeReplicaContext(peer);
        RegionStoreClient client = clientBuilder.build(region, triple.getMiddle());

        if (future == null) {
            future = new CompletableFuture<>();
        }
        CompletableFuture<Kvrpcpb.RedisResponse> finalFuture = future;

        client.execRedisAsync(context, backOffer, executor, sendTime, cmd, encodedKey, value)
            .whenComplete((response, throwable) -> {
                if (response != null) {
                    finalFuture.complete(response);
                    return;
                }
                if (throwable instanceof RetryableException) {
                    long time = ((RetryableException) throwable).getRetryTime();
                    executor.schedule(() -> {
                            Triple<TiRegion, TiStore, Metapb.Peer> target =
                                tiSession.getRegionManager().getRegionStoreByKey(encodedKey, backOffer, false);
                            directExecRedis(executor, finalFuture, target, backOffer, sendTime, cmd, encodedKey, value);
                        },
                        time, TimeUnit.MILLISECONDS);
                    return;
                }
                finalFuture.completeExceptionally(throwable);
            });

        return finalFuture;
    }

    public void close() {
        if (tiSession.isClosed()) {
            return;
        }
        tiSession.close();
    }

}
